
#include "glodef.h"
#include "grid3d.h"
#include "errlog.h"
#include "matrices.h"
#include "fe.h"
#include "fe_matr.h"
#include "poisson.h"

#include <stdlib.h>

#include "ns_pcorr.h"

/* note: the grid variable used in this file comes from glovar.h */
#include "glovar.h"


static vector AV_pc;      /* global variable, which is local to this file */



typedef struct
{
    unsigned id;
    unsigned eltype[2];
    unsigned ndofs[2];
    vector r;
    vector t;
    scalar d;
    scalar q;
    void *poisson;
    fe_matr_t *momentum;
} pcorr_t;


#define PCORR_ID   0x406033F0

static pcorr_t *get_pdata(void *data)
{
    pcorr_t *p = data;
    if((data==NULL)||(p->id!=PCORR_ID))
        error_msg("ns_pcorr.c: invalid poisson_t structure passed.\n");
    return p;
}


static pcorr_t *_check_input(void *pdata,
        const vector *u,  const scalar *phi,  const scalar *p,
        const vector *u0, const scalar *phi0, const scalar *p0)
{
    pcorr_t *pd =get_pdata(pdata);
    if( u==NULL || u->N!=pd->ndofs[0] )
        goto _error_;
    if( u0==NULL || u0->N!=pd->ndofs[0] )
        goto _error_;
    if( phi==NULL || phi->N!=pd->ndofs[1] )
        goto _error_;
    if( phi0==NULL || phi0->N!=pd->ndofs[1] )
        goto _error_;
    if( p==NULL || p->N!=pd->ndofs[1] )
        goto _error_;
    if( p0==NULL || p0->N!=pd->ndofs[1] )
        goto _error_;
    return pd; /* happily */
    _error_:
    error_msg("Invalid input to TimeStep_pc.\n");
}


/****************************************************************************
 *                                                                          *
 *                                                                          *
 *                                                                          *
 ****************************************************************************
 */


void TimeStep_pc(void *pdata, vector *u, scalar *phi, scalar *p,
        const vector *adv_vel,
        const vector *u0, const scalar *phi0, const scalar *p0,
        unsigned reassemble)
{
    pcorr_t *pd = _check_input(pdata, u, phi, p, u0, phi0, p0);
    unsigned i;
    
    PrintLog("TimeStep using pressure correction scheme.\n");
    
    if(flags.time_order != 1)
        error_msg("Only 1st order scheme is available.\n");
    
    /************************************************************************
     *          STEP 1: The momentum equation
     */
    
    /*
     * Prepare the right-hand-side of the system
     */
    clean_vec(&pd->r);
    /* the source term */
    test_vector_field(grid, pd->eltype[0], TIME, srce_trm, &pd->t);
    for(i=0;i<3*pd->ndofs[0];i++)
        pd->r.X[i]+=TStep*pd->t.X[i];
    /* the prevoius velocity */
    test_vector(grid, pd->eltype[0], u0, &pd->t);
    for(i=0;i<3*pd->ndofs[0];i++)
        pd->r.X[i]+=pd->t.X[i];
    /* the pressure gradient */
    for(i=0;i<pd->ndofs[1];i++)
        pd->q.F[i]=p0->F[i]+phi0->F[i];
    /*
     * test_grad_e1_in_e2(grid,pd->eltype[1],&pd->q,pd->eltype[0],&pd->t);
     * for(i=0;i<3*pd->ndofs[0];i++)
     * pd->r.X[i]-=TStep*pd->t.X[i];
     */
    test_e1_in_div_e2(grid, pd->eltype[1], &pd->q, pd->eltype[0], &pd->t);
    for(i=0;i<3*pd->ndofs[0];i++)
        pd->r.X[i]+=TStep*pd->t.X[i];
    
    /*
     * Handle advection
     */
    if(adv_vel == NULL)
        adv_vel = u0;
    switch(flags.advection)
    {
        case adv_none: /* do nothing */
            break;
        case adv_explicit: /* explicit advection -- modify pd->r */
            clean_vec(&AV_pc);   /* keep matrix as is */
            test_advection_vector(grid, pd->eltype[0], adv_vel, u0, &pd->t);
            for(i=0;i<3*pd->ndofs[0];i++)
                pd->r.X[i]-=TStep*pd->t.X[i];
            break;
        case adv_semi: /* semi-implicit advection -- modify matrix */
            copy_vec(adv_vel, &AV_pc);
            reassemble = TRUE;      /* must reassemble matrix */
            break;
        case adv_char: /* characteristics : do nothing */
            break;
        default:
            error_msg("Illegal advection code %d\n", flags.advection);
    }
    
    /*
     * Apply boundary conditions
     */
    apply_bc_dir_vector(grid, pd->eltype[0], TIME, bdry_vel, &pd->r);
    
    if(on_ival | on_oval)
        div_free_bc(grid, pd->eltype[0], on_ival | on_oval, &pd->r);
    
    for(i=0;i<pd->ndofs[0];i++)
        if(fe_dof_bcond(grid, pd->eltype[0], i)==BC_DIR)
        {
            u->X[i] = pd->r.X[i];
            u->Y[i] = pd->r.Y[i];
            u->Z[i] = pd->r.Z[i];
        }
    
    /*
     * Solve the momentum equation
     */
    if(reassemble)
    {
        fe_matr_assemble(grid, 1, &pd->momentum);
    }
    PrintLog("Solving the momentum equations.\n");
    /* inital guess for u is u itself (assuming u has the values of u0 */
    fe_matr_solve_vector(pd->momentum, u, &pd->r, &pd->t, 0.0);
    
    /************************************************************************
     *          STEP 2: The Poisson problem
     */
    test_div_e1_in_e2(grid, pd->eltype[0], u, pd->eltype[1], &pd->d);
    for(i=0;i<pd->ndofs[1];i++)
        pd->q.F[i] = pd->d.F[i]*TStep;
    /* What is a good initial guess for phi? phi=0? */
    clean_scal(phi);
    SolvePoisson(pd->poisson, phi, &pd->q, &pd->d, TRUE, reassemble);
    
    /***********************************************************************
     *          STEP 3: The pressure correction
     */
    /* test phi+p0 and put in q */
    for(i=0;i<pd->ndofs[1];i++)
        pd->d.F[i] = phi->F[i]+p0->F[i];
    test_scalar(grid, pd->eltype[1], &pd->d, &pd->q);
    /* test div.u in the pressure element */
    test_div_e1_in_e2(grid, pd->eltype[0], u, pd->eltype[1], &pd->d);
    {
        FLOAT Coef =  -1.0/GetRe();
        for(i=0;i<pd->ndofs[1];i++)
            pd->q.F[i] += Coef*pd->d.F[i];
    }
    copy_scal(p0, p); /* initial guess */
    /* invert the mass matrix */
    PrintLog("Inverting the pressure mass matrix.\n");
    fe_matr_solve_scalar(massP, p, &pd->q, &pd->d, 0.0);
}



static void loc_moment(const grid3d_t *grid, unsigned el, unsigned type, FLOAT ***S)
{
    static FLOAT _M[MAX_LDOFS][MAX_LDOFS];
    static FLOAT *M[MAX_LDOFS];
    FLOAT mass_coef, stif_coef, advm_coef;
    /* extern vector AV_pc; */
    FLOAT **lM, **lS;
    FLOAT **lA, **lD;
    FLOAT Rho = 1.0;
    FLOAT Mu  = 1.0 / GetRe();
    unsigned nn, i, j;
    for(i=0;i<MAX_LDOFS;i++)
        M[i]=_M[i];
    if(S)
        *S=M;
    else
        return;
    get_coefs(flags.advection, &mass_coef, &stif_coef, &advm_coef);
    local_mass(grid, el, type, &lM);
    local_stif(grid, el, type, &lS);
    nn=fe_dofs(grid, el, type, NULL);
    for(i=0;i<nn;i++)
        for(j=0;j<nn;j++)
            (*S)[i][j] = Rho*mass_coef*lM[i][j]+Mu*stif_coef*lS[i][j];
    if(flags.advection==adv_semi)
    {
        local_advm(grid, el, type, &AV_pc, &lA);
        local_advm_D(grid, el, type, &AV_pc, &lD);
        for(i=0;i<nn;i++)
            for(j=0;j<nn;j++)
                (*S)[i][j] += Rho*advm_coef*(lA[i][j]+lD[i][j]);
    }
}

void *InitPCorr(unsigned elemU, unsigned elemP)
{
    pcorr_t *pd;
    int symmetric;
    
    InitFE(1, &elemU, grid);
    PrintLog("Initialize pressure correction scheme ...");
    
    pd=malloc(sizeof(pcorr_t));
    if(pd==NULL)
        error_msg("ns_pcorr.c: memory allocation error.\n");
    assert(grid!=NULL);
    pd->id=PCORR_ID;
    pd->eltype[0]=elemU;
    pd->eltype[1]=elemP;
    pd->ndofs[0] = fe_ndof(grid, elemU);
    pd->ndofs[1] = fe_ndof(grid, elemP);
    alloc_vec(&AV_pc, pd->ndofs[0]);
    alloc_vec(&pd->r, pd->ndofs[0]);
    alloc_vec(&pd->t, pd->ndofs[0]);
    alloc_scal(&pd->q, pd->ndofs[1]);
    alloc_scal(&pd->d, pd->ndofs[1]);
    if(flags.advection==adv_semi)
        symmetric = FALSE;
    else 
        symmetric = TRUE;
    pd->momentum=fe_matr_alloc(elemU, loc_moment, eo_bc_dir, grid, symmetric, 
            NULL);
    fe_matr_assemble(grid, 1, &pd->momentum);
    
    if(flags.need_adjoint)
        error_msg("Adjoint solver is not available with pressure correction.\n");
    
    PrintLog("\tsuccess.\n");
    
    pd->poisson=InitPoisson(grid, elemP);
    return pd;
}

void DonePCorr(void *pdata)
{
    pcorr_t *pd = get_pdata(pdata);
    DonePoisson(pd->poisson);
    free_vec(&AV_pc);
    free_vec(&pd->r);
    free_vec(&pd->t);
    free_scal(&pd->q);
    free_scal(&pd->d);
    fe_matr_free(pd->momentum);
    free(pd);
    
    PrintLog("Pressure correction scheme done.\n");
}

